# main.py

import logging
import config
from data_loader import load_finetuning_datasets, prefix_dataset_ids
from finetuner import FinetuningExperiment

# --- Setup Logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')


def main():
    """Main execution script for finetuning experiments."""
    logger = logging.getLogger(__name__)

    # --- 1. Initial Experiment on Scifact ---
    logger.info(f"--- Starting Initial Experiment on '{config.INITIAL_DATASET_NAME}' ---")

    # Load data
    train_dataset, val_dataset, test_dataset = load_finetuning_datasets(config.INITIAL_DATASET_NAME)

    # Initialize experiment runner
    experiment = FinetuningExperiment(config)

    # Run experiments
    baseline_score = experiment.run_baseline(train_dataset, test_dataset)
    adapter_score = experiment.run_adapter_finetuning(train_dataset, test_dataset)
    nudge_score = experiment.run_nudge_finetuning(train_dataset, val_dataset, test_dataset)

    # Display results
    print("\n" + "=" * 50)
    print("           INITIAL EXPERIMENT RESULTS")
    print("=" * 50)
    print(f"Dataset: {config.INITIAL_DATASET_NAME}")
    print(f"Metric: NDCG@{config.RETRIEVER_TOP_K}\n")
    print(f"Baseline (BGE-small):          {baseline_score:.4f}")
    print(f"Adapter Finetuning:            {adapter_score:.4f}")
    print(f"NUDGE Finetuning:              {nudge_score:.4f}")
    print("=" * 50 + "\n")

    # --- 2. Data Insertion Experiment with NFCorpus ---
    logger.info(f"--- Starting Data Insertion Experiment with '{config.INSERTION_DATASET_NAME}' ---")

    # Load and prepare new data
    new_train, new_val, new_test = load_finetuning_datasets(config.INSERTION_DATASET_NAME)
    new_train = prefix_dataset_ids(new_train, config.INSERTION_ID_PREFIX)
    new_val = prefix_dataset_ids(new_val, config.INSERTION_ID_PREFIX)
    new_test = prefix_dataset_ids(new_test, config.INSERTION_ID_PREFIX)
    logger.info(f"Prefixed all IDs in '{config.INSERTION_DATASET_NAME}' dataset with '{config.INSERTION_ID_PREFIX}-'")

    # Run insertion experiment
    score_new, score_original_after_insert = experiment.run_nudge_insertion(
        new_train_dataset=new_train,
        new_val_dataset=new_val,
        new_test_dataset=new_test,
        original_test_dataset=test_dataset  # Pass the original scifact test set
    )

    # Display insertion results
    print("\n" + "=" * 50)
    print("           DATA INSERTION EXPERIMENT RESULTS (NUDGE)")
    print("=" * 50)
    print("Evaluated on Aggregated Corpus (scifact + nfcorpus)\n")
    print(f"Performance on new data ({config.INSERTION_DATASET_NAME}):   {score_new:.4f}")
    print(
        f"Performance on original data ({config.INITIAL_DATASET_NAME}): {score_original_after_insert:.4f} (No catastrophic forgetting)")
    print("=" * 50 + "\n")


if __name__ == "__main__":
    main()